conda create -n EAST python=3.7
conda activate EAST
pip install shapely
pip install opencv-python==4.0.0.21
pip install lanms-proper

python3 train.py

sudo apt update
sudo apt install libsm6

from src.util import AverageMeter, get_param_groups
from src.east import EAST, EastWithLossCell
from src.logger import get_logger
from src.initializer import default_recurisive_init
from src.dataset import create_east_dataset
from src.lr_scheduler import get_lr

args, _ = parser.parse_known_args()
args.device_id = int(os.getenv("DEVICE_ID", "0"))
args.rank = args.device_id
 
ms.set_context(mode=ms.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id)
if args.is_distributed:
    comm.init()
    args.rank = comm.get_rank()
    args.group_size = comm.get_group_size()
    ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.DATA_PARALLEL, gradients_mean=True,
                                 device_num=args.group_size)

if args.is_modelArts:
    import moxing as mox
 
    local_data_url = os.path.join('/cache/data', str(args.rank))
    local_ckpt_url = os.path.join('/cache/ckpt', str(args.rank))
    local_ckpt_url = os.path.join(local_ckpt_url, 'backbone.ckpt')
 
    mox.file.rename(args.pretrained_backbone, local_ckpt_url)
    args.pretrained_backbone = local_ckpt_url
 
    mox.file.copy_parallel(args.data_dir, local_data_url)
    args.data_dir = local_data_url
 
    args.outputs_dir = os.path.join('/cache', args.outputs_dir)

args.data_root = os.path.abspath(os.path.join(args.data_dir, 'image'))
args.txt_root = os.path.abspath(os.path.join(args.data_dir, 'groundTruth'))
 
# 使用当前进程的编号（即 args.rank 变量）作为子目录名称，拼接成完整的输出文件夹路径
outputs_dir = os.path.join(args.outputs_dir, str(args.rank))
# 获取当前时间作为子目录名称，再次拼接成完整的输出文件夹路径
args.outputs_dir = os.path.join(
    args.outputs_dir,
    datetime.datetime.now().strftime('%Y-%m-%d_time_%H_%M_%S'))
args.logger = get_logger(args.outputs_dir, args.rank)  # 调用 get_logger()函数创建一个日志记录器，并将日志保存在 args.outputs_dir 目录下
args.logger.save_args(args)  # 将所有参数保存在日志文件中

if args.need_profiler:
    # 创建一个性能分析器，并将结果保存在args.outputs_dir路径下
    profiler = Profiler(
        output_path=args.outputs_dir,
        is_detail=True,
        is_show_op_path=True)
 
# 创建一个AverageMeter对象用于记录损失值的平均值，以便后续输出和打印
loss_meter = AverageMeter('loss')
 
# 重置自动并行上下文
context.reset_auto_parallel_context()
parallel_mode = ParallelMode.STAND_ALONE
degree = 1
# 又是分布式计算相关……
if args.is_distributed:
    parallel_mode = ParallelMode.DATA_PARALLEL
    degree = args.group_size
context.set_auto_parallel_context(
    parallel_mode=parallel_mode,
    gradients_mean=True,
    device_num=degree)

network = EAST()  # 设置 network，加载 EAST 模型
# default is kaiming-normal
default_recurisive_init(network)  # 对 EAST 模型进行默认的递归初始化。这里使用的是 kaiming-normal（He 正态分布）初始化方法
 
# load pretrained_backbone
if args.pretrained_backbone:  # 如果不为 None，载入预训练的 backbone 模型
    parm_dict = load_checkpoint(args.pretrained_backbone)  # 加载模型参数
    load_param_into_net(network, parm_dict)  # 将模型参数加载到 network 上
    args.logger.info('finish load pretrained_backbone')  # 在日志中记录加载完成的信息
 
network = EastWithLossCell(network)  # 将 EAST 模型和损失函数进行结合，即将模型传入 EastWithLossCell()函数，得到组合后的模型对象
if args.resume_east:  # 如果 args.resume_east 不为 None，继续训练之前保存的 EAST 模型，resume：恢复，继续
    param_dict = load_checkpoint(args.resume_east)
    load_param_into_net(network, param_dict)
    args.logger.info('finish get resume east')
 
args.logger.info('finish get network')
 
# 载入数据集，调用 create_east_dataset()函数，传入图片文件夹路径、文本文件夹路径、批量大小、设备数量、进程编号等参数，获取数据集以及数据总数，并在日志中记录加载完成的信息。
ds, data_size = create_east_dataset(img_root=args.data_root, txt_root=args.txt_root, batch_size=args.per_batch_size,
                                    device_num=args.group_size, rank=args.rank, is_training=True)
args.logger.info('Finish loading dataset')
 
# 计算每个 epoch 中的步数，即将数据总数、批量大小和设备数量进行计算得到
args.steps_per_epoch = int(
    data_size /
    args.per_batch_size /
    args.group_size)
 
if not args.ckpt_interval:
    # 如果 args.ckpt_interval 为空，则将其设置为每个 epoch 的步数
    args.ckpt_interval = args.steps_per_epoch

# get learnning rate
lr = get_lr(args)  # 函数获取当前epoch的学习率，并将其赋值给变量lr
opt = Adam(  # 使用Adam优化器进行优化，并指定优化器的参数为EAST模型中需要更新的参数
    params=get_param_groups(network),
    learning_rate=Tensor(
        lr,
        ms.float32))
loss_scale = FixedLossScaleManager(1.0, drop_overflow_update=True)  # 固定的损失缩放管理器
model = Model(network, optimizer=opt, loss_scale_manager=loss_scale)  # 使用Model函数从EAST模型对象和优化器拼接出一个完整的训练模型，并将损失缩放管理器传入
# 这样就生成了完整的训练模型对象，并且可以对其进行训练

network.set_train()  # 将网络设置为训练状态
# save the network model and parameters for subsequence fine-tuning
# 设置保存检查点的配置信息，包括保存检查点的步数和最大保存数量，并将其赋值给变量 config_ck
config_ck = CheckpointConfig(
    save_checkpoint_steps=100,
    keep_checkpoint_max=1)
# group layers into an object with training and evaluation features
# 指定模型参数保存路径
save_ckpt_path = os.path.join(
    args.outputs_dir, 'ckpt_' + str(args.rank) + '/')
# 使用 ModelCheckpoint()函数创建一个回调函数，用于保存训练模型参数
ckpoint_cb = ModelCheckpoint(
    prefix="checkpoint_east",
    directory=save_ckpt_path,
    config=config_ck)
# 创建一个回调函数，用于保存训练模型参数。其中，prefix 参数指定保存文件名的前缀，directory 参数指定保存路径，config 参数指定保存配置信息。
callback = [
    TimeMonitor(data_size=data_size),
    LossMonitor(),
    ckpoint_cb
]
# 调用 model.train()方法对训练模型进行训练，传入总 epoch 数、数据集以及之前定义的回调函数列表。在训练过程中，启用了数据集下沉模式，即 dataset_sink_mode=True，以提高训练效率
model.train(
    args.max_epoch,
    ds,
    callbacks=callback,
    dataset_sink_mode=True)
args.logger.info('========end training=============')

class EAST(nn.Cell):
    def __init__(self):
        super(EAST, self).__init__()
        # 提取图像特征的模块，返回 5 组特征图用于后续处理
        self.extractor = VGG16FeatureExtraction()
        # 将特征图组合的模块，将 5 组特征图拼接在一起，形成更为丰富多样的特征信息用于后续处理
        self.merge = Merge()
        # 输出模块，对拼接后的特征图进行卷积处理来得到文本区域预测分数 score 和几何信息预测值 geo
        self.output = Output()
 
    def construct(self, x_1):
        # 通过 x_1 输入数据调用 self.extractor()获取 5 组特征图
        f_0, f_1, f_2, f_3, f_4 = self.extractor(x_1)
        # 将这些特征图传入 self.merge()模块进行拼接，得到拼接后的特征图
        x_1 = self.merge(f_0, f_1, f_2, f_3, f_4)
        # 将该特征图输入到 self.output()模块获得文本区域预测分数 score 和几何信息预测值 geo
        score, geo = self.output(x_1)
 
        return score, geo

class VGG16FeatureExtraction(nn.Cell):
    """VGG16FeatureExtraction for deeptext"""
 
    def __init__(self):
        super(VGG16FeatureExtraction, self).__init__()
        self.relu = nn.ReLU()
        self.max_pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.avg_pool = nn.AvgPool2d(kernel_size=2, stride=2)
 
        self.conv1_1 = _conv(
            in_channels=3,
            out_channels=64,
            kernel_size=3,
            padding=1)
 
        ……
        
        self.conv5_3 = _conv(
            in_channels=512,
            out_channels=512,
            kernel_size=3,
            padding=1)
        self.cast = P.Cast()
 
    def construct(self, out):
        """ Construction of VGG """
        f_0 = out
        out = self.cast(out, mstype.float32)
        out = self.conv1_1(out)
        out = self.relu(out)
        out = self.conv1_2(out)
        out = self.relu(out)
        out = self.max_pool(out)
 
		……
 
        out = self.max_pool(out)
        f_4 = out
        out = self.conv5_1(out)
        out = self.relu(out)
        out = self.conv5_2(out)
        out = self.relu(out)
        out = self.conv5_3(out)
        out = self.relu(out)
        out = self.max_pool(out)
        f_5 = out
 
        return f_0, f_2, f_3, f_4, f_5

class Merge(nn.Cell):
    def __init__(self):
        super(Merge, self).__init__()
 
        self.conv1 = nn.Conv2d(1024, 128, 1, has_bias=True)
        self.bn1 = nn.BatchNorm2d(128)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(
            128,
            128,
            3,
            padding=1,
            pad_mode='pad',
            has_bias=True)
        self.bn2 = nn.BatchNorm2d(128)
        self.relu2 = nn.ReLU()
 
		……
 
    def construct(self, x, f1, f2, f3, f4):
        img_hight = P.Shape()(x)[2]
        img_width = P.Shape()(x)[3]
 
        out = P.ResizeBilinear((img_hight / 16, img_width / 16), True)(f4)
        out = self.concat((out, f3))
        out = self.relu1(self.bn1(self.conv1(out)))
        out = self.relu2(self.bn2(self.conv2(out)))
 
        out = P.ResizeBilinear((img_hight / 8, img_width / 8), True)(out)
        out = self.concat((out, f2))
        out = self.relu3(self.bn3(self.conv3(out)))
        out = self.relu4(self.bn4(self.conv4(out)))
 
        out = P.ResizeBilinear((img_hight / 4, img_width / 4), True)(out)
        out = self.concat((out, f1))
        out = self.relu5(self.bn5(self.conv5(out)))
        out = self.relu6(self.bn6(self.conv6(out)))
 
        out = self.relu7(self.bn7(self.conv7(out)))
        return out

class Output(nn.Cell):
    def __init__(self, scope=512):
        super(Output, self).__init__()
        self.conv1 = nn.Conv2d(32, 1, 1)
        self.sigmoid1 = nn.Sigmoid()
        self.conv2 = nn.Conv2d(32, 4, 1)
        self.sigmoid2 = nn.Sigmoid()
        self.conv3 = nn.Conv2d(32, 1, 1)
        self.sigmoid3 = nn.Sigmoid()
        self.scope = scope
        self.concat = P.Concat(axis=1)
        self.PI = 3.1415926535898
 
    def construct(self, x):
        score = self.sigmoid1(self.conv1(x))  # 文本区域得分
        loc = self.sigmoid2(self.conv2(x)) * self.scope  # 位置
        angle = (self.sigmoid3(self.conv3(x)) - 0.5) * self.PI  # 倾斜角度
        geo = self.concat((loc, angle))  # 边界框信息包含位置和倾斜角度
        return score, geo  # 最终返回文本区域得分和拼接后的边界框信息

class EastLossBlock(nn.Cell):
    def __init__(self):
        super(EastLossBlock, self).__init__()
        self.split = P.Split(1, 5)
        self.min = MyMin()
        self.log = P.Log()
        self.cos = P.Cos()
        self.mean = P.ReduceMean(keep_dims=False)
        self.sum = P.ReduceSum()
        self.eps = 1e-5
        self.dice = DiceCoefficient()
 
    def construct(
            self,
            y_true_cls,
            y_pred_cls,
            y_true_geo,
            y_pred_geo,
            training_mask):
        ans = self.sum(y_true_cls)
        # 将预测得到的文本区域得分与真实标注的文本区域得分进行比较，采用 Dice 系数计算分类损失
        classification_loss = self.dice(
            y_true_cls, y_pred_cls * (1 - training_mask))
 
        # n * 5 * h * w
        # 将预测得到的位置信息和真实标注的位置信息分别拆分出来
        d1_gt, d2_gt, d3_gt, d4_gt, theta_gt = self.split(y_true_geo)
        d1_pred, d2_pred, d3_pred, d4_pred, theta_pred = self.split(y_pred_geo)
        area_gt = (d1_gt + d3_gt) * (d2_gt + d4_gt)
        area_pred = (d1_pred + d3_pred) * (d2_pred + d4_pred)
        w_union = self.min(d2_gt, d2_pred) + self.min(d4_gt, d4_pred)
        h_union = self.min(d1_gt, d1_pred) + self.min(d3_gt, d3_pred)
 
        area_intersect = w_union * h_union
        area_union = area_gt + area_pred - area_intersect
        # 通过计算交并比（IoU）和角度误差得到位置损失
        iou_loss_map = -self.log((area_intersect + 1.0) /
                                 (area_union + 1.0))  # iou_loss_map
        angle_loss_map = 1 - self.cos(theta_pred - theta_gt)  # angle_loss_map
		
        # 角度误差使用余弦相似度计算
        angle_loss = self.sum(angle_loss_map * y_true_cls) / ans
        iou_loss = self.sum(iou_loss_map * y_true_cls) / ans
        geo_loss = 10 * angle_loss + iou_loss
 
        return geo_loss + classification_loss

class EastWithLossCell(nn.Cell):
    def __init__(self, network):
        super(EastWithLossCell, self).__init__()
        # 传入一个EAST模型，作为计算图中的网络模块
        self.east_network = network
        # 实例化了EastLossBlock类，作为计算图中的损失函数模块
        self.loss = EastLossBlock()
 
    def construct(self, img, true_cls, true_geo, training_mask):
        '''
        img: 输入图片
        true_cls: 分类标注
        true_geo: 位置标注
        training_mask: 训练集中的掩码（用于过滤掉非真实文本区域的数据）
        '''
        # 调用计算图进行前向计算
        socre, geometry = self.east_network(img)
        # 将计算得到的分类得分和位置信息分别传给损失函数模块进行后向计算，得到整体的损失值并返回
        loss = self.loss(
            true_cls,
            socre,
            true_geo,
            geometry,
            training_mask)
        return loss

def create_east_dataset(
        img_root,
        txt_root,
        batch_size,
        device_num,
        rank,
        is_training=True):
    # 实例化 ICDAREASTDataset 类，传入图片和文本标注的路径，用于读取并解析图像和标注
    east_data = ICDAREASTDataset(img_path=img_root, gt_path=txt_root)
    # 生成分布式采样器，用于在多个设备之间对数据进行划分和分发。
    distributed_sampler = DistributedSampler(
        len(east_data), device_num, 0 if device_num == 1 else rank, shuffle=True)
 
    trans_list = [CV.RandomColorAdjust(0.5, 0.5, 0.5, 0.25),  # 随机改变图像的颜色饱和度、对比度和亮度
                  CV.Rescale(1 / 255.0, 0),  # 对图像进行缩放
                  CV.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # 图像正则化处理
                  CV.HWC2CHW()]  # 将图像的通道维度从 HWC（高×宽×通道数）顺序转换为 CHW（通道数×高×宽）顺序
    if is_training:  # 如果是训练模式
        dataset_column_names = [
            "image",  # 图像
            "score_map",  # 分类标注分数图
            "geo_map",  # 位置标注几何图
            "training_mask"]  # 训练集掩码
        # 调用 MindSpore 中的 GeneratorDataset 类生成数据集
        ds = de.GeneratorDataset(
            east_data,
            column_names=dataset_column_names,
            num_parallel_workers=32,  # 数据处理和增强过程中使用的并行线程数
            # sampler 参数则指定了数据采样器，即从数据集中选择数据样本的方式，
            # 本例中使用的是前面提到的分布式采样器 distributed_sampler
            sampler=distributed_sampler)
        # 调用 map()方法将数据集中的图像列传入变换列表中的操作进行增广
        ds = ds.map(
            operations=trans_list,
            input_columns=["image"],
            num_parallel_workers=8,
            python_multiprocessing=True)
        # 使用 batch()方法将批量大小对数据集进行划分
        ds = ds.batch(batch_size, num_parallel_workers=8, drop_remainder=True)
 
    return ds, len(east_data)

class ICDAREASTDataset:
    def __init__(self, img_path, gt_path, scale=0.25, length=512):
        super(ICDAREASTDataset, self).__init__()
        self.img_files = [os.path.join(
            img_path,
            img_file) for img_file in sorted(os.listdir(img_path))]
        self.gt_files = [
            os.path.join(
                gt_path,
                gt_file) for gt_file in sorted(
                    os.listdir(gt_path))]
        self.scale = scale  # 缩放比例
        self.length = length  # 裁剪后的图像长度
 
    def __getitem__(self, index):
        with open(self.gt_files[index], 'r') as f:
            lines = f.readlines()
        vertices, labels = extract_vertices(lines)  # 从文本标注中提取文本区域的顶点坐标和标注
 
        img = Image.open(self.img_files[index])  # 读取图像
        img, vertices = adjust_height(img, vertices)  # 调整高度
        img, vertices = rotate_img(img, vertices)  # 随机旋转图像
        img, vertices = crop_img(img, vertices, labels, self.length)  # 将图像切割成指定长度的大小
        score_map, geo_map, ignored_map = get_score_geo(
            img, vertices, labels, self.scale, self.length)  # 分类标注分数图、位置标注几何图和忽略标注
        score_map = score_map.transpose(2, 0, 1)
        ignored_map = ignored_map.transpose(2, 0, 1)
        geo_map = geo_map.transpose(2, 0, 1)
        if np.sum(score_map) < 1:
            score_map[0, 0, 0] = 1
        return img, score_map, geo_map, ignored_map
 
    def __len__(self):
        return len(self.img_files)

def extract_vertices(lines):
    '''extract vertices info from txt lines
    Input:
      lines   : list of string info 输入是一个字符串列表 lines，其中每个字符串包含了一个文本区域的信息，包括顶点坐标和标签等
    Output:
      vertices: vertices of text regions <numpy.ndarray, (n,8)>  所有文本区域的顶点坐标
      labels  : 1->valid, 0->ignore, <numpy.ndarray, (n,)>  标签
    '''
    labels = []  # 存储最终的标签
    vertices = []  # 存储顶点信息
    for line in lines:
        # 通过 rstrip()和 lstrip()函数去除其前后空格和 BOM（Byte Order Mark）等特殊字符，并使用 split()函数将其切分为一个包含八个整数的列表
        vertices.append(list(map(int, line.rstrip('\n').lstrip('\ufeff').split(',')[:8])))
        label = 0 if '###' in line else 1
        labels.append(label)
    # 返回顶点和标签的 numpy 数组
    return np.array(vertices), np.array(labels)

def adjust_height(img, vertices, ratio=0.2):
    '''adjust height of image to aug data
    Input:
      img         : PIL Image
      vertices    : vertices of text regions <numpy.ndarray, (n,8)>
      ratio       : height changes in [0.8, 1.2]
    Output:
      img         : adjusted PIL Image
      new_vertices: adjusted vertices
    '''
    ratio_h = 1 + ratio * (np.random.rand() * 2 - 1)  # 随机调整输入图像的高度
    old_h = img.height
    # 根据输入的高度缩放比例ratio_h，计算调整后的图像新高度new_h。
    # 原始图像的高度由变量old_h指定，通过乘以缩放比例并四舍五入取整来得到调整后的高度。
    # np.around()函数是NumPy库中的一个函数，用于对数组进行四舍五入，其默认精度为0
    new_h = int(np.around(old_h * ratio_h))
    img = img.resize((img.width, new_h), Image.BILINEAR)
 
    new_vertices = vertices.copy()
    if vertices.size > 0:
        new_vertices[:, [1, 3, 5, 7]] = vertices[:, [1, 3, 5, 7]] * (new_h / old_h)
    # 返回调整后的图像和更新后的顶点坐标
    return img, new_vertices

def rotate_img(img, vertices, angle_range=10):
    '''rotate image [-10, 10] degree to aug data
    Input:
      img         : PIL Image
      vertices    : vertices of text regions <numpy.ndarray, (n,8)>
      angle_range : rotate range
    Output:
      img         : rotated PIL Image
      new_vertices: rotated vertices
    '''
    # 获得中心旋转点
    center_x = (img.width - 1) / 2
    center_y = (img.height - 1) / 2
    angle = angle_range * (np.random.rand() * 2 - 1)
    # 使用了 BILINEAR 滤波器来进行图像插值，以获得更好的旋转效果
    img = img.rotate(angle, Image.BILINEAR)
    # 定义一个大小为 vertices.shape 的全零 NumPy 数组 new_vertices，用于存储旋转后的顶点坐标
    new_vertices = np.zeros(vertices.shape)
    for i, vertice in enumerate(vertices):
        # 遍历每个文本区域的顶点坐标，调用 rotate_vertices()函数来计算旋转后的新坐标，然后将其保存到 new_vertices 中
        new_vertices[i, :] = rotate_vertices(
            vertice, -angle / 180 * math.pi, np.array([[center_x], [center_y]]))
    return img, new_vertices

def crop_img(img, vertices, labels, length):
    '''crop img patches to obtain batch and augment
    Input:
      img         : PIL Image
      vertices    : vertices of text regions <numpy.ndarray, (n,8)>
      labels      : 1->valid, 0->ignore, <numpy.ndarray, (n,)>
      length      : length of cropped image region
    Output:
      region      : cropped image region
      new_vertices: new vertices in cropped region
    '''
    # 获取原始图像的高度h和宽度w
    h, w = img.height, img.width
    # confirm the shortest side of image >= length
    # 如果其中较小的一边小于指定的裁剪长度，则使用PIL库提供的resize()方法将图像缩放到相应的大小
    if h >= w and w < length:
        img = img.resize((length, int(h * length / w)), Image.BILINEAR)
    elif h < w and h < length:
        img = img.resize((int(w * length / h), length), Image.BILINEAR)
    ratio_w = img.width / w
    ratio_h = img.height / h
    assert (ratio_w >= 1 and ratio_h >= 1)
 
    # 如果其中较小的一边小于指定的裁剪长度，则使用PIL库提供的resize()方法将图像缩放到相应的大小
    new_vertices = np.zeros(vertices.shape)
    if vertices.size > 0:
        new_vertices[:, [0, 2, 4, 6]] = vertices[:, [0, 2, 4, 6]] * ratio_w
        new_vertices[:, [1, 3, 5, 7]] = vertices[:, [1, 3, 5, 7]] * ratio_h
 
    # find random position
    # 生成随机的裁剪位置，检查裁剪区域是否与文本区域相交，避免将裁剪区域中的文本区域遮盖或截断
    remain_h = img.height - length
    remain_w = img.width - length
    flag = True
    cnt = 0
    while flag and cnt < 1000:
        # 若随机裁剪的位置与文本区域有交集，则继续生成新的随机位置，
        # 直到找到一个合适的位置或者超过最大尝试次数1000次为止
        cnt += 1
        start_w = int(np.random.rand() * remain_w)
        start_h = int(np.random.rand() * remain_h)
        flag = is_cross_text([start_w, start_h], length,
                             new_vertices[labels == 1, :])
    box = (start_w, start_h, start_w + length, start_h + length)
    # 使用PIL库提供的crop()方法从原始图像中截取指定大小的区域，并将其作为本函数的输出返回。
    region = img.crop(box)
    if new_vertices.size == 0:
        # 如果不存在任何文本区域，则直接返回裁剪后的图像区域和空的新顶点坐标
        return region, new_vertices
    # 更新文本区域的顶点坐标。将新的裁剪图像左上角的坐标(start_w, start_h)作为原点，计算相对于这个原点的顶点坐标，并将这个相对坐标赋值给new_vertices
    new_vertices[:, [0, 2, 4, 6]] -= start_w
    new_vertices[:, [1, 3, 5, 7]] -= start_h
    return region, new_vertices

context.set_context(
    mode=context.GRAPH_MODE,  # 图模式
    device_target=args.device_target,  # 设备类型
    save_graphs=False,  # 是否保存计算图
    device_id=args.device_num)  # 设备编号

if __name__ == '__main__':
    model_name = args.checkpoint_path
    test_img_path = args.test_img_path
    submit_path = './submit'
    eval_model(model_name, test_img_path, submit_path)

def eval_model(name, img_path, submit, save_flag=True):
    '''
    name: 模型的 checkpoint 文件路径
    img_path: 测试集图片所在的文件夹路径
    submit: 输出结果保存的文件夹路径
    save_flag: 是否保存中间结果
    '''
    # 判断输出结果保存的目录是否存在，如果存在则删除该目录及其子目录，然后重新创建一个同名目录
    if os.path.exists(submit):
        shutil.rmtree(submit)
    os.mkdir(submit)
    # 构建 EAST 模型
    network = EAST()
    # 加载预训练权重参数
    param_dict = load_checkpoint(name)
    load_param_into_net(network, param_dict)
    # 设置模型为训练模式
    network.set_train(True)
 
    start_time = time.time()
    # 调用 detect_dataset()函数对测试集图片进行检测，并将检测结果保存到指定的输出目录 submit 中
    detect_dataset(network, img_path, submit)
    os.chdir(submit)
    res = subprocess.getoutput('zip -q submit.zip *.txt')
    res = subprocess.getoutput('mv submit.zip ../')
    os.chdir('../')
    # 调用评估脚本./evaluate/script.py 来计算模型的性能指标，评估结果保存在字符串变量 res 中
    res = subprocess.getoutput(
        'python ./evaluate/script.py -g=./evaluate/gt.zip -s=./submit.zip')
    print(res)
    os.remove('./submit.zip')
    print('eval time is {}'.format(time.time() - start_time))
 
    if not save_flag:
        # 如果 save_flag 为 False，则删除输出目录及其子目录（闻到了屎山的味道）
        shutil.rmtree(submit)

def detect_dataset(model, test_img_path, submit_path):
    """
    detection on whole dataset, save .txt results in submit_path
        Input:
                model        : detection model 模型实例
                device       : gpu if gpu is available
                test_img_path: dataset path 测试图片所在文件夹的路径
                submit_path  : submit result for evaluation 提交结果保存路径
    """
    # 读取测试集中所有的图片，并按照文件名排序
    img_files = os.listdir(test_img_path)
    img_files = sorted([os.path.join(test_img_path, img_file)
                        for img_file in img_files])
 
    for i, img_file in enumerate(img_files):
        # 对于每一张图片，调用detect()函数进行目标检测，返回目标框的坐标信息
        print('evaluating {} image'.format(i), end='\r')
        boxes = detect(Image.open(img_file), model)
        seq = []
        if boxes is not None:
            # 如果检测结果不为空，则将框的坐标信息转换成符合要求的字符串序列并加入到列表seq中
            seq.extend([','.join([str(int(b))
                                  for b in box[:-1]]) + '\n' for box in boxes])
        # 将序列seq保存为与当前图片名称相同的.txt文件格式，并将其写入submit_path目录下
        with open(os.path.join(submit_path, 'res_' +
                               os.path.basename(img_file).replace('.jpg', '.txt')), 'w') as f:
            f.writelines(seq)  # 当检测完成后，输出log信息提示检测进度

def detect(img, model):
    """detect text regions of img using model
        Input:
                img   : PIL Image
                model : detection model
                device: gpu if gpu is available
        Output:
                detected polys
    """
    # 将输入图片进行尺寸调整与相应的 ratio 变换，得到调整后的图片、高宽比例 ratio_h 和 ratio_w
    img, ratio_h, ratio_w = resize_img(img)
    # 利用模型对调整后的图片进行文字区域检测，得到概率图 score 和文本框参数 geo
    score, geo = model(load_pil(img))
    # 对概率图和文本框参数使用 PaddlePaddle 中的 Squeeze()函数进行维度降低（由 4 维转为 3 维）
    score = P.Squeeze(0)(score)
    geo = P.Squeeze(0)(geo)
    # 从降维后的概率图和文本框参数中获取文本框坐标信息，即调用 get_boxes()函数
    boxes = get_boxes(score.asnumpy(), geo.asnumpy())
    # 根据之前的高宽比例 ratio_h 和 ratio_w，调整并计算出检测到的文本框在原始图片上的坐标信息，即调用 adjust_ratio()函数
    return adjust_ratio(boxes, ratio_w, ratio_h)

def get_boxes(score, geo, score_thresh=0.9, nms_thresh=0.2):
    """get boxes from feature map
        Input:
                score       : score map from model <numpy.ndarray, (1,row,col)> 概率图
                geo         : geo map from model <numpy.ndarray, (5,row,col)> 文本框参数
                score_thresh: threshold to segment score map 置信度阈值
                nms_thresh  : threshold in nms 非极大值抑制阈值
        Output:
                boxes       : final polys <numpy.ndarray, (n,9)>
        """
    # 对输入的score进行降维，即将其转化为二维数组
    score = score[0, :, :]
    # 在降维后的score数组中，找到大于score_thresh的点，并以(r,c)的格式记录下来，形成一个n x 2的矩阵xy_text
    xy_text = np.argwhere(score > score_thresh)  # n x 2, format is [r, c]
    # 按行排序xy_text，以保证前面的点在结果中优先考虑
    if xy_text.size == 0:
        return None
 
    # 将xy_text中的坐标信息转化为正确的x,y坐标（由于降维之前是按行major的顺序排列，因此需要将列号作为x坐标，行号作为y坐标）
    xy_text = xy_text[np.argsort(xy_text[:, 0])]
    valid_pos = xy_text[:, ::-1].copy()  # n x 2, [x, y]
    # 从降维后的geo数组中提取出与xy_text中相应位置点相关的文本框参数，形成5 x n的矩阵valid_geo
    valid_geo = geo[:, xy_text[:, 0], xy_text[:, 1]]  # 5 x n
    # 利用restore_polys()函数将valid_pos和valid_geo还原为文本框的坐标点集polys_restored，并得到对应的索引值index
    polys_restored, index = restore_polys(valid_pos, valid_geo, score.shape)
    if polys_restored.size == 0:
        return None
 
    # 将polys_restored表示为(n,8)大小的数组，其中前8列分别为文本框像素点的坐标，第9列为该文本框的置信度（即所在score map中的值）
    boxes = np.zeros((polys_restored.shape[0], 9), dtype=np.float32)
    boxes[:, :8] = polys_restored
    boxes[:, 8] = score[xy_text[index, 0], xy_text[index, 1]]
    # 对polys_restored执行非极大值抑制（NMS）操作，得到最终的文本框坐标信息boxes
    boxes = lanms.merge_quadrangle_n9(boxes.astype('float32'), nms_thresh)
    return boxes

def adjust_ratio(boxes, ratio_w, ratio_h):
    """refine boxes
        Input:
                boxes  : detected polys <numpy.ndarray, (n,9)>
                ratio_w: ratio of width
                ratio_h: ratio of height
        Output:
                refined boxes
        """
    if boxes is None or boxes.size == 0:
        return None
    boxes[:, [0, 2, 4, 6]] /= ratio_w
    boxes[:, [1, 3, 5, 7]] /= ratio_h
    return np.around(boxes)

source activate base  # 第一次进服务器激活需要 activate base
python -c "import mindspore;mindspore.run_check()"  # 查看 mindspore 版本
conda create -n east --clone base  # 克隆 base 环境
conda activate east  # 激活 east 环境
pip install numpy
pip install opencv-python
pip install shapely
pip install pillow
pip install lanms-neo
pip install --upgrade setuptools  # 更新 setuptools
pip install Polygon3  # 这个库很难装，可能需要更新 setuptools
pip install onnxruntime

cd /home/ma-user/work/east/
python3 train.py

python3 eval.py

Calculated!{"precision": 0.527431421446384, "recall": 0.6109773712084737, "hmean": 0.566138746375195, "AP": 0}